package com.unisound.tracelog.client.core.context;

import org.apache.commons.collections4.CollectionUtils;

import java.util.ArrayList;
import java.util.List;

public class SpanContextManager {
    /**
     * 保持线程隔离(不支持父子线程)
     * 子线程不能直接感知父线程的数据,需要参数直接传递,避免受到父线程数据更新的影响
     */
    private static final ThreadLocal<List<SpanContext>> spanContextThreadLocal = new ThreadLocal<>();

    public static void add(SpanContext spanContext) {
        if (!check(spanContext)) return;

        SpanContext cloneSpanContext = spanContext.clone();
        if (cloneSpanContext == null) return;

        List<SpanContext> spanContexts = getSpanContexts();

        if (spanContexts == null) {
            spanContexts = new ArrayList<>();
            spanContexts.add(cloneSpanContext);
            securitySet(spanContexts);
        } else if (spanContexts.size() == 0) {
            spanContexts.add(cloneSpanContext);
            securitySet(spanContexts);
        } else {
            SpanContext lastSpanContext = spanContexts.get(spanContexts.size() - 1);
            if (!cloneSpanContext.equals(lastSpanContext)) {
                spanContexts.add(cloneSpanContext);
                securitySet(spanContexts);
            }
        }
    }

    private static void securitySet(List<SpanContext> spanContexts) {
        if (CollectionUtils.isEmpty(spanContexts)) return;
        List<SpanContext> newSpanContexts = new ArrayList<>();
        for (SpanContext spanContext : spanContexts) {
            add(spanContext, newSpanContexts);
        }

        if (CollectionUtils.isNotEmpty(newSpanContexts)) {
            spanContextThreadLocal.set(newSpanContexts);
        }
    }

    private static void add(SpanContext spanContext, List<SpanContext> spanContexts) {
        if (!check(spanContext)) return;
        SpanContext cloneSpanContext = spanContext.clone();
        if (spanContexts == null || cloneSpanContext == null) return;
        spanContexts.add(cloneSpanContext);
    }

    public static void update(List<SpanContext> spanContexts) {
        securitySet(spanContexts);
    }

    private static boolean check(SpanContext spanContext) {
        if (spanContext == null ||
                spanContext.getTraceId() == null ||
                spanContext.getSpanId() == null ||
                spanContext.getParentSpanId() == null) return false;
        return true;
    }

    public static void remove() {
        List<SpanContext> spanContexts = getSpanContexts();
        if (CollectionUtils.isEmpty(spanContexts)) return;

        spanContexts.remove(spanContexts.size() - 1);

        if (spanContexts.isEmpty()) {
            spanContextThreadLocal.remove();
        } else {
            securitySet(spanContexts);
        }
    }

    public static List<SpanContext> getSpanContexts() {
        List<SpanContext> spanContexts = spanContextThreadLocal.get();
        if (CollectionUtils.isEmpty(spanContexts)) return null;

        List<SpanContext> newSpanContexts = new ArrayList<>();
        for (SpanContext spanContext : spanContexts) {
            SpanContext cloneSpanContext = spanContext.clone();
            if (cloneSpanContext != null) {
                newSpanContexts.add(cloneSpanContext);
            }
        }
        return newSpanContexts;
    }

    public static SpanContext getSpanContext() {
        List<SpanContext> spanContexts = getSpanContexts();
        if (CollectionUtils.isEmpty(spanContexts)) return null;
        SpanContext spanContext = spanContexts.get(spanContexts.size() - 1);
        return spanContext;
    }

    public static boolean isRoot() {
        return getSpanContext() == null;
    }

}
